# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：drawfigs.py
    Created on：2025/4/9 22:36
    Author：(input)
    Description: 绘图代码
"""
import numpy as np
import matplotlib.pyplot as plt

class SosiFlowPlotter:
    """
    SosiFlowPlotter 用于绘制源汇流的物理现象，如流函数、势函数、速度场等。

    该类接收 PhysicalGrid 提供的网格信息和 SourceSink 提供的流动参数。
    """

    def __init__(self, phys_grid, aux_grid, source_sink):
        """
        初始化 SosiFlowPlotter

        Parameters:
        ----------
        phys_grid : PhysicalGrid
            提供物理平面网格信息的 PhysicalGrid 对象
        aux_grid : AuxiliaryGrid
            提供辅助平面网格信息的 AuxiliaryGrid 对象
        source_sink : SourceSink
            包含源和汇流动参数的 SourceSink 对象
        """
        # 存储网格和源汇对象
        self.phys_grid = phys_grid
        self.aux_grid = aux_grid
        self.source_sink = source_sink

        # 获取物理平面网格坐标
        self.X = self.phys_grid.X
        self.Y = self.phys_grid.Y
        self.Z = self.phys_grid.to_complex()

        # 获取辅助平面网格坐标
        self.U = self.aux_grid.U
        self.V = self.aux_grid.V
        self.W = self.aux_grid.to_complex()

        # 计算绘图范围
        self.x_range = (self.phys_grid.x.min(), self.phys_grid.x.max())
        self.y_range = (self.phys_grid.y.min(), self.phys_grid.y.max())
        self.u_range = (self.aux_grid.U.min(), self.aux_grid.U.max())
        self.v_range = (self.aux_grid.V.min(), self.aux_grid.V.max())

        # 计算宽高比
        self.aspect_ratio = (self.phys_grid.x[1] - self.phys_grid.x[0]) / (self.phys_grid.y[1] - self.phys_grid.y[0])

        # 计算统一的 fig_size
        base_width = 10  # 基础宽度（英寸）
        base_height = base_width / self.aspect_ratio  # 根据宽高比计算高度
        self.fig_size = (base_width, base_height)

        # 物理平面上的源汇坐标
        self.source_x = np.real(self.source_sink.z_so)
        self.source_y = np.imag(self.source_sink.z_so)
        self.sink_x = np.real(self.source_sink.z_si)
        self.sink_y = np.imag(self.source_sink.z_si)

        # 辅助平面上的源汇坐标
        self.w_source_x = np.real(self.source_sink.w_so)
        self.w_source_y = np.imag(self.source_sink.w_so)
        self.w_sink_x = np.real(self.source_sink.w_si)
        self.w_sink_y = np.imag(self.source_sink.w_si)

        self.contour_set = None

    def plot_real_geo(self, ax=None, is_physical=True):
        """
        绘制源汇流的物理环境，包括源点、汇点和边界。

        Parameters:
        ----------
        ax : matplotlib.axes.Axes, optional
            绘图的轴对象，如果为 None 则使用当前轴
        is_physical : bool
            是否绘制物理平面（True）或辅助平面（False）
        """
        if ax is None:
            ax = plt.gca()

        # 源、汇位置
        if is_physical:
            ax.scatter(self.source_x, self.source_y, color='r', s=80, marker='o', label='Source (Inlet)')
            ax.scatter(self.sink_x, self.sink_y, color='g', s=80, marker='X', label='Sink (Outlet)')

            # 绘制矩形边界
            wall_left = self.phys_grid.x.min()
            wall_right = self.phys_grid.x.max()
            wall_bottom = self.phys_grid.y.min()
            wall_top = self.phys_grid.y.max()
            ax.plot([wall_left, wall_right], [wall_top, wall_top], 'k-', lw=1)  # 上边界
            ax.plot([wall_left, wall_right], [wall_bottom, wall_bottom], 'k-', lw=1)  # 下边界
            ax.plot([wall_left, wall_left], [wall_bottom, wall_top], 'k-', lw=1)  # 左边界
            ax.plot([wall_right, wall_right], [wall_bottom, wall_top], 'k-', lw=1)  # 右边界
        else:
            ax.scatter(self.w_source_x, self.w_source_y, color='r', s=80, marker='o', label='Source (Inlet)')
            ax.scatter(self.w_sink_x, self.w_sink_y, color='g', s=80, marker='X', label='Sink (Outlet)')
            # 绘制辅助平面的实轴（矩形边界映射到实轴）
            ax.plot([self.u_range[0], self.u_range[1]], [0, 0], 'k-', lw=1, label='Mapped Boundary (Real Axis)')

        ax.legend(loc='upper right')

    def plot_contour(self, data, levels, cmap, colorbar_label, title, xlabel='Width', ylabel='Height', num_levels=12, is_physical=True):
        """
        通用等值线绘制方法，适用于绘制流函数 (psi)、势函数 (phi) 和其他场信息。

        Parameters:
        ----------
        data : np.ndarray
            需要绘制的二维数据（如势函数phi或流函数psi）
        levels : list, optional
            等值线的级别
        cmap : str
            颜色映射（colormap）
        colorbar_label : str
            色条标签
        title : str
            图像标题
        xlabel : str
            x 轴标签
        ylabel : str
            y 轴标签
        num_levels : int
            等值线数量
        is_physical : bool
            是否绘制物理平面（True）或辅助平面（False）
        """
        # 检查数据类型
        if not np.issubdtype(data.dtype, np.number):
            raise ValueError(f"{colorbar_label} must contain numeric values")
        min_val, max_val = np.min(data), np.max(data)
        if min_val == max_val:
            raise ValueError(f"{colorbar_label} must have a range of values for contouring")

        # 动态调整等值线数量
        if levels is None:
            range_val = max_val - min_val
            num_levels = min(20, max(5, int(range_val / (np.nanmean(np.abs(data)) + 1e-10) * 10)))
            levels = np.linspace(min_val, max_val, num_levels)

        plt.figure(figsize=self.fig_size)
        ax = plt.gca()

        levels = levels if levels is not None else np.linspace(min_val, max_val, num_levels)
        contour = plt.contour(self.X if is_physical else self.U, self.Y if is_physical else self.V, data, levels=levels,
                              cmap=cmap, linewidths=2)
        plt.clabel(contour, inline=True, fontsize=10, fmt="%.2e")
        plt.colorbar(contour, label=colorbar_label)

        self.contour_set = contour  # 存储等值线数据

        # 设置绘图范围
        if is_physical:
            plt.xlim(self.x_range[0], self.x_range[1])
            plt.ylim(self.y_range[0], self.y_range[1])
        else:
            plt.xlim(self.u_range[0], self.u_range[1])
            plt.ylim(self.v_range[0], self.v_range[1])

        self.plot_real_geo(ax, is_physical)

        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        plt.grid(True)
        plt.axis("equal")
        plt.show()

    def extract_contour_coordinates(self, contour_value):
        """
        从等值线中提取指定值的坐标点。

        Parameters:
        ----------
        contour_value : float
            要提取的等值线水平值

        Returns:
        -------
        coords : list
            包含该等值线所有坐标点的列表
        """
        if not hasattr(self, 'contour_set'):
            raise AttributeError("No contour set available. Run a contour plot method first.")
        coords = []
        for collection in self.contour_set.collections:
            for path in collection.get_paths():
                vertices = path.vertices
                if np.isclose(contour_value, collection.get_array()).any():
                    coords.append(vertices)
        return coords

    def plot_psi(self, psi, title='Stream Function', is_physical=True):
        """绘制流函数的等值线图"""
        self.plot_contour(psi, levels=None,
                          cmap="coolwarm",
                          colorbar_label='Stream Function',
                          title=title,
                          xlabel='Width' if is_physical else 'u',
                          ylabel='Height' if is_physical else 'v',
                          is_physical=is_physical)

    def plot_phi(self, phi, title='Potential Function', is_physical=True):
        """绘制势函数的等值线图"""
        self.plot_contour(phi, levels=None,
                          cmap="coolwarm",
                          colorbar_label='Potential Function',
                          title=title,
                          xlabel='Width' if is_physical else 'u',
                          ylabel='Height' if is_physical else 'v',
                          is_physical=is_physical)

    def plot_flow_net(self, phi, psi, title='Flow Net', num_levels=12, is_physical=True):
        """
        绘制流网图，展示势函数和流函数的等值线。
        """
        # 检查数据
        if not (np.issubdtype(phi.dtype, np.number) and np.issubdtype(psi.dtype, np.number)):
            raise ValueError("phi and psi must contain numeric values")
        min_phi, max_phi = np.min(phi), np.max(phi)
        min_psi, max_psi = np.min(psi), np.max(psi)
        if min_phi == max_phi or min_psi == max_psi:
            raise ValueError("phi and psi must have a range of values for contouring")

        plt.figure(figsize=self.fig_size)
        ax = plt.gca()

        # 绘制网格点作为背景
        #if is_physical:
            #plt.scatter(self.X.flatten(), self.Y.flatten(), c='gray', s=5, alpha=0.3)
        #else:
            #plt.scatter(self.U.flatten(), self.V.flatten(), c='gray', s=5, alpha=0.3)

        levels_phi = np.linspace(min_phi, max_phi, num_levels)
        levels_psi = np.linspace(min_psi, max_psi, num_levels)

        contour_phi = plt.contour(self.X if is_physical else self.U, self.Y if is_physical else self.V, phi,
                                  levels=levels_phi, cmap="viridis", linewidths=2)
        plt.clabel(contour_phi, inline=True, fontsize=10, fmt="%.2e")
        contour_psi = plt.contour(self.X if is_physical else self.U, self.Y if is_physical else self.V, psi,
                                  levels=levels_psi, cmap="coolwarm", linewidths=2)
        plt.clabel(contour_psi, inline=True, fontsize=10, fmt="%.2e")
        self.contour_set = contour_psi  # 存储最后一个等值线集合

        self.plot_real_geo(ax, is_physical)

        if is_physical:
            plt.xlim(self.x_range[0], self.x_range[1])
            plt.ylim(self.y_range[0], self.y_range[1])
        else:
            plt.xlim(self.u_range[0], self.u_range[1])
            plt.ylim(self.v_range[0], self.v_range[1])

        plt.xlabel('Width' if is_physical else 'u')
        plt.ylabel('Height' if is_physical else 'v')
        plt.title(title)
        plt.grid(True)
        plt.axis("equal")
        plt.show()

    def plot_velocity_field(self, vx, vy, magV=None, show_magnitude=False, title='Velocity Field',
                            cmap='jet', is_physical=True, clip_percentile=99):
        """
        绘制速度场流线图，用箭头表示方向，可选通过颜色和颜色条表示速度大小。

        Parameters:
        ----------
        vx : np.ndarray
            水平速度分量
        vy : np.ndarray
            垂直速度分量
        magV : np.ndarray, optional
            速度模（大小），当 show_magnitude=True 时必须提供
        show_magnitude : bool, optional
            是否显示速度大小，默认 False
        title : str, optional
            图表标题，默认 'Velocity Field'
        cmap : str, optional
            颜色映射，默认 'jet'
        is_physical : bool
            是否绘制物理平面（True）或辅助平面（False）
        clip_percentile : float, optional
            对速度大小进行裁剪的百分位数（用于剔除奇点影响），默认 99
        """
        # 检查数据
        if not (np.issubdtype(vx.dtype, np.number) and np.issubdtype(vy.dtype, np.number)):
            raise ValueError("vx and vy must contain numeric values")
        if np.min(vx) == np.max(vx) and np.min(vy) == np.max(vy):
            raise ValueError("vx and vy must have a range of values for stream plotting")

        plt.figure(figsize=self.fig_size)
        ax = plt.gca()

        if show_magnitude:
            if magV is None:
                raise ValueError("magV must be provided when show_magnitude is True")
            if not np.issubdtype(magV.dtype, np.number):
                raise ValueError("magV must contain numeric values")
            if np.min(magV) == np.max(magV):
                raise ValueError("magV must have a range of values for color mapping")

            # 裁剪速度，防止奇点影响色阶
            v_clip = np.percentile(magV, clip_percentile)
            magV_clipped = np.clip(magV, 0, v_clip)

            stream = plt.streamplot(self.X if is_physical else self.U,
                                    self.Y if is_physical else self.V,
                                    vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->',
                                    color=magV_clipped, cmap=cmap)
            plt.colorbar(stream.lines, label=f'Velocity Magnitude (clipped at {clip_percentile}%)')
        else:
            stream = plt.streamplot(self.X if is_physical else self.U,
                                    self.Y if is_physical else self.V,
                                    vx, vy,
                                    density=1, linewidth=1,
                                    arrowsize=1, arrowstyle='->')

        self.plot_real_geo(ax, is_physical)

        if is_physical:
            plt.xlim(self.x_range[0], self.x_range[1])
            plt.ylim(self.y_range[0], self.y_range[1])
        else:
            plt.xlim(self.u_range[0], self.u_range[1])
            plt.ylim(self.v_range[0], self.v_range[1])

        plt.xlabel('Width' if is_physical else 'u')
        plt.ylabel('Height' if is_physical else 'v')
        plt.title(title)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.axis("equal")
        plt.tight_layout()
        plt.show()

    def plot_pressure(self, p, title='Pressure Coefficient Distribution', num_levels=20, is_physical=True):
        """
        绘制压力分布等值线图。
        """
        # 检查数据
        if not np.issubdtype(p.dtype, np.number):
            raise ValueError("Pressure must contain numeric values")
        min_p, max_p = np.min(p), np.max(p)
        if min_p == max_p:
            raise ValueError("Pressure must have a range of values for contouring")

        plt.figure(figsize=self.fig_size)
        ax = plt.gca()

        # 绘制网格点作为背景
        if is_physical:
            plt.scatter(self.X.flatten(), self.Y.flatten(), c='gray', s=5, alpha=0.3)
        else:
            plt.scatter(self.U.flatten(), self.V.flatten(), c='gray', s=5, alpha=0.3)

        contour = plt.contour(self.X if is_physical else self.U, self.Y if is_physical else self.V, p,
                              levels=num_levels, cmap='viridis', linewidths=1)
        self.contour_set = contour
        plt.clabel(contour, inline=True, fontsize=10, fmt="%.2e")

        self.plot_real_geo(ax, is_physical)

        plt.colorbar(contour, label='Pressure Coefficient (Cp)', extend="both")

        if is_physical:
            plt.xlim(self.x_range[0], self.x_range[1])
            plt.ylim(self.y_range[0], self.y_range[1])
        else:
            plt.xlim(self.u_range[0], self.u_range[1])
            plt.ylim(self.v_range[0], self.v_range[1])

        plt.xlabel('Width' if is_physical else 'u')
        plt.ylabel('Height' if is_physical else 'v')
        plt.title(title)
        plt.grid(True)
        plt.axis("equal")
        plt.show()